// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/core/extension_type.h>

#include <memory>
#include <mutex>
#include <sstream>
#include <string>
#include <unordered_map>
#include <utility>

#include <nebula/array/util.h>
#include <nebula/core/chunked_array.h>
#include <nebula/core/config.h>
#include <nebula/extension/fixed_shape_tensor.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <turbo/log/logging.h>

namespace nebula {


    DataTypeLayout ExtensionType::layout() const { return storage_type_->layout(); }

    std::string ExtensionType::to_string(bool show_metadata) const {
        std::stringstream ss;
        ss << "extension<" << this->extension_name() << ">";
        return ss.str();
    }

    std::shared_ptr<Array> ExtensionType::WrapArray(const std::shared_ptr<DataType> &type,
                                                    const std::shared_ptr<Array> &storage) {
                DKCHECK_EQ(type->id(), Type::EXTENSION);
        const auto &ext_type = turbo::checked_cast<const ExtensionType &>(*type);
                DKCHECK_EQ(storage->type_id(), ext_type.storage_type()->id());
        auto data = storage->data()->copy();
        data->type = type;
        return ext_type.make_array(std::move(data));
    }

    std::shared_ptr<ChunkedArray> ExtensionType::WrapArray(
            const std::shared_ptr<DataType> &type, const std::shared_ptr<ChunkedArray> &storage) {
                DKCHECK_EQ(type->id(), Type::EXTENSION);
        const auto &ext_type = turbo::checked_cast<const ExtensionType &>(*type);
                DKCHECK_EQ(storage->type()->id(), ext_type.storage_type()->id());

        ArrayVector out_chunks(storage->num_chunks());
        for (int i = 0; i < storage->num_chunks(); i++) {
            auto data = storage->chunk(i)->data()->copy();
            data->type = type;
            out_chunks[i] = ext_type.make_array(std::move(data));
        }
        return std::make_shared<ChunkedArray>(std::move(out_chunks));
    }

    ExtensionArray::ExtensionArray(const std::shared_ptr<ArrayData> &data) { set_data(data); }

    ExtensionArray::ExtensionArray(const std::shared_ptr<DataType> &type,
                                   const std::shared_ptr<Array> &storage) {
        KCHECK_EQ(type->id(), Type::EXTENSION);
        KCHECK(
                storage->type()->equals(*turbo::checked_cast<const ExtensionType &>(*type).storage_type()));
        auto data = storage->data()->copy();
        // XXX This pointer is reverted below in set_data()...
        data->type = type;
        set_data(data);
    }

    void ExtensionArray::set_data(const std::shared_ptr<ArrayData> &data) {
        KCHECK_EQ(data->type->id(), Type::EXTENSION);
        this->Array::set_data(data);

        auto storage_data = data->copy();
        storage_data->type = (static_cast<const ExtensionType &>(*data->type).storage_type());
        storage_ = make_array(storage_data);
    }

    class ExtensionTypeRegistryImpl : public ExtensionTypeRegistry {
    public:
        ExtensionTypeRegistryImpl() {}

        turbo::Status register_type(std::shared_ptr<ExtensionType> type) override {
            std::lock_guard<std::mutex> lock(lock_);
            std::string type_name = type->extension_name();
            auto it = name_to_type_.find(type_name);
            if (it != name_to_type_.end()) {
                return turbo::already_exists_error("A type extension with name ", type_name,
                                               " already defined");
            }
            name_to_type_[type_name] = std::move(type);
            return turbo::OkStatus();
        }

        turbo::Status UnregisterType(const std::string &type_name) override {
            std::lock_guard<std::mutex> lock(lock_);
            auto it = name_to_type_.find(type_name);
            if (it == name_to_type_.end()) {
                return turbo::not_found_error("No type extension with name ", type_name, " found");
            }
            name_to_type_.erase(it);
            return turbo::OkStatus();
        }

        std::shared_ptr<ExtensionType> get_type(const std::string &type_name) override {
            std::lock_guard<std::mutex> lock(lock_);
            auto it = name_to_type_.find(type_name);
            if (it == name_to_type_.end()) {
                return nullptr;
            } else {
                return it->second;
            }
            return nullptr;
        }

    private:
        std::mutex lock_;
        std::unordered_map<std::string, std::shared_ptr<ExtensionType>> name_to_type_;
    };

    static std::shared_ptr<ExtensionTypeRegistry> g_registry;
    static std::once_flag registry_initialized;

    namespace internal {

        static void CreateGlobalRegistry() {
            g_registry = std::make_shared<ExtensionTypeRegistryImpl>();

            // Register canonical extension types
            auto ext_type =
                    turbo::checked_pointer_cast<ExtensionType>(extension::fixed_shape_tensor(int64(), {}));

            KCHECK_OK(g_registry->register_type(ext_type));
        }

    }  // namespace internal

    std::shared_ptr<ExtensionTypeRegistry> ExtensionTypeRegistry::GetGlobalRegistry() {
        std::call_once(registry_initialized, internal::CreateGlobalRegistry);
        return g_registry;
    }

    turbo::Status RegisterExtensionType(std::shared_ptr<ExtensionType> type) {
        auto registry = ExtensionTypeRegistry::GetGlobalRegistry();
        return registry->register_type(type);
    }

    turbo::Status UnregisterExtensionType(const std::string &type_name) {
        auto registry = ExtensionTypeRegistry::GetGlobalRegistry();
        return registry->UnregisterType(type_name);
    }

    std::shared_ptr<ExtensionType> GetExtensionType(const std::string &type_name) {
        auto registry = ExtensionTypeRegistry::GetGlobalRegistry();
        return registry->get_type(type_name);
    }

    extern const char kExtensionTypeKeyName[] = "ARROW:extension:name";
    extern const char kExtensionMetadataKeyName[] = "ARROW:extension:metadata";

}  // namespace nebula
